// Copyright (c) 1996-2002 John Lyon-Smith. All rights reserved.

#pragma once

#include <wxtl\Base.h>


#pragma warning(push)
#pragma warning(disable:4239)	// non-standard extension

static const TCHAR *HCU = _T("HKEY_CURRENT_USER");
static const TCHAR *HLM = _T("HKEY_LOCAL_MACHINE");

namespace wxtl
{
	struct REG_KEY_INFO
	{
		tstring strClass;			// class name 
		DWORD nSubKeys;				// number of subkeys 
		DWORD nMaxSubkey;			// longest subkey name length  
		DWORD nMaxClass;			// longest class string length 
		DWORD nValues;				// number of value entries 
		DWORD nMaxValueName;		// longest value name length 
		DWORD nMaxValueData;		// longest value data length 
		DWORD pbSecurityDescriptor;	// security descriptor length 
		FILETIME ftLastWriteTime;	// last write time 
	};

	// This class doesn't inherit from Handle because sadly HANDLE and HKEY are not the same type

	class RegistryKey
	{
	public:
		// Valid base keys are:
		//
		// HKEY_CURRENT_USER
		// HKEY_LOCAL_MACHINE
		// HKEY_CLASSES_ROOT
		// HKEY_USERS
		// HKEY_PERFORMANCE_DATA
		// HKEY_CURRENT_CONFIG
		// HKEY_DYN_DATA
		//
		// We don't define an enum because we need the HKEY type for method overloading

		RegistryKey() : m_hkey(NULL) 
		{ 
		}
		RegistryKey(HKEY hKey) 
		{ 
			m_hkey = hKey; 
		}
		RegistryKey(RegistryKey& key) throw() :
			m_hkey(NULL)
		{
			attach(key.detach());
		}

		~RegistryKey() throw()
		{
			close();
		}

		RegistryKey& operator=(RegistryKey& key) throw()
		{
			if (this != &key)
			{
				attach(key.detach());
			}

			return *this;
		}

		HKEY get() const throw()
		{
			return m_hkey;
		}

		void attach(HKEY h) throw()
		{
			close();
			m_hkey = h;
		}

		HKEY detach() throw()
		{
			HKEY h;

			h = m_hkey;
			m_hkey = NULL;

			return h;
		}

		virtual bool isValid() const
		{
			return (m_hkey != NULL);
		}

		void close() throw(_win32_error)
		{
			if (isValid())
			{
				::RegCloseKey(m_hkey);
			}
		}

		RegistryKey createKey(	
			PCTSTR subKeyName,
			TCHAR *className = NULL,
			DWORD dwOptions = REG_OPTION_NON_VOLATILE,
			REGSAM sam = KEY_ALL_ACCESS,
			DWORD* pdwDisp = NULL) const throw(_win32_error)
		{
			HKEY hKey;
			DWORD dwDisp;
			
			THROW_BAD_WINERROR_RET(::RegCreateKeyEx(
				m_hkey,	
				subKeyName,
				0,
				className,
				dwOptions,
				sam,
				NULL,
				&hKey,
				&dwDisp));

			if (NULL != pdwDisp)
				*pdwDisp = dwDisp;

			return RegistryKey(hKey);
		}

		RegistryKey openKey(PCTSTR subKey, REGSAM sam = KEY_ALL_ACCESS)
			const throw(_win32_error)
		{
			HKEY hKey;
			
			THROW_BAD_WINERROR_RET(::RegOpenKeyEx(m_hkey, subKey, 0, sam, &hKey));

			return RegistryKey(hKey);
		}

		void deleteKey(PCTSTR szSubKey) throw(_win32_error)
		{
			THROW_BAD_WINERROR_RET(::RegDeleteKey(m_hkey, szSubKey));
		}

		void deleteValue(PCTSTR szValue) throw(_win32_error)
		{
			THROW_BAD_WINERROR_RET(::RegDeleteValue(m_hkey, szValue));
		}

		bool enumKey(
			DWORD& iSubkey, 
			tstring& strName, 
			tstring& strClass, 
			FILETIME &ftLastWrite) throw(_win32_error)
		{
			TCHAR keyName[256];
			DWORD keyNameSize = sizeof(keyName)/sizeof(TCHAR);
			TCHAR className[256];
			DWORD classNameSize = sizeof(className)/sizeof(TCHAR);

			DWORD error = ::RegEnumKeyEx(
				m_hkey,
				iSubkey,
				keyName,
				&keyNameSize,
				NULL,
				className,
				&classNameSize,
				&ftLastWrite);

			strName = keyName;
			strClass = className;
			
			if (ERROR_SUCCESS == error)
			{
				iSubkey++;
				return true;
			}
			else if (ERROR_NO_MORE_ITEMS == error)
			{
				return false;
			}
			else
			{
				throw _win32_error(error);
				return false;
			}
		}

		bool enumValue(
			DWORD& iIndex, 
			tstring& strName,
			DWORD& dwValueType, 
			LPBYTE pData, 
			DWORD dwDataSize,
			DWORD& dwActualSize) throw(_win32_error)
		{
			TCHAR keyName[256];
			DWORD keyNameSize = sizeof(keyName)/sizeof(TCHAR);
			
			dwActualSize = dwDataSize;
			
			DWORD error = ::RegEnumValue(
				m_hkey,
				iIndex,
				keyName,
				&keyNameSize,
				NULL,
				&dwValueType,
				pData,
				&dwActualSize);

			if (ERROR_SUCCESS == error)
			{
				strName = keyName;
				iIndex++;
				return true;
			}
			else if (ERROR_NO_MORE_ITEMS == error)
			{
				return false;
			}
			else
				throw _win32_error(error);
		}

		void flush() throw(_win32_error)
		{
			THROW_BAD_WINERROR_RET(::RegFlushKey(m_hkey));
		}

		void queryKeyInfo(REG_KEY_INFO& rki) throw(_win32_error)
		{
			TCHAR className[256];
			DWORD classNameSize = sizeof(className)/sizeof(TCHAR);

			THROW_BAD_WINERROR_RET(::RegQueryInfoKey(
				m_hkey,
				className,
				&classNameSize,
				NULL,
				&rki.nSubKeys,
				&rki.nMaxSubkey,
				&rki.nMaxClass,
				&rki.nValues,
				&rki.nMaxValueName,
				&rki.nMaxValueData,
				&rki.pbSecurityDescriptor,
				&rki.ftLastWriteTime));
		}

		void queryValueInfo(
			PCTSTR pszValueName, 
			DWORD* pdwValueType = NULL,
			DWORD* pdwValueByteSize = NULL) throw(_win32_error)
		{
			THROW_BAD_WINERROR_RET(::RegQueryValueEx(
				m_hkey,
				pszValueName,
				NULL,
				pdwValueType,
				NULL,
				pdwValueByteSize));
		}

		void queryValue(PCTSTR pszValueName) throw(_win32_error)
		{
			return queryValueInfo(pszValueName);
		}

		void queryValue(PCTSTR pszValueName, PVOID pData, DWORD& dwDataLen) throw(_win32_error)
		{
			DWORD nValueType = 0;
			bool moreData;
			
			queryValue(pszValueName, nValueType, (LPBYTE)pData, dwDataLen, moreData); 

			if (REG_BINARY != nValueType)
				throw _win32_error(ERROR_INVALID_FUNCTION);
		}

		void queryValue(PCTSTR pszValueName, tstring& strValue, bool bAutoExpand = true) throw(_win32_error)
		{
			DWORD dwSize = 1024;
			TCHAR *pValue = new TCHAR[dwSize + 1];
			DWORD nValueType = 0;
			bool moreData;
			
			queryValue(pszValueName, nValueType, (LPBYTE)pValue, dwSize, moreData);

			if (REG_SZ != nValueType && REG_EXPAND_SZ != nValueType)
				throw _win32_error(ERROR_INVALID_FUNCTION);

			// Some keys just keep growing...
			while (moreData)
			{
				delete [] pValue;
				pValue = new TCHAR[dwSize + 1];
				
				queryValue(pszValueName, nValueType, (LPBYTE)pValue, dwSize, moreData);
			}

			if (REG_EXPAND_SZ == nValueType && bAutoExpand)
			{
				TCHAR* pNewValue = new TCHAR[dwSize];
				DWORD dwNewSize = ::ExpandEnvironmentStrings(pValue, pNewValue, dwSize);
				
				if (dwNewSize > dwSize)
				{
					delete [] pNewValue;
					
					pNewValue = new TCHAR[dwNewSize];
					
					::ExpandEnvironmentStrings(pNewValue, pValue, dwNewSize);
				}

				delete [] pValue;
				pValue = pNewValue;
			}

			strValue = pValue;
			delete [] pValue;
		}

		void queryValueShort(PCTSTR pszValueName, tstring& strValue, bool bAutoExpand = true) throw(_win32_error)
		{
			DWORD dwSize = 1024;
			//TCHAR *pValue = new TCHAR[dwSize + 1];
			TCHAR pValue[2000] = {0};
			DWORD nValueType = 0;
			bool moreData;

			queryValue(pszValueName, nValueType, (LPBYTE)pValue, dwSize, moreData);

			if (REG_SZ != nValueType && REG_EXPAND_SZ != nValueType)
				throw _win32_error(ERROR_INVALID_FUNCTION);

			// Some keys just keep growing...
			/*while (moreData)
			{
				delete [] pValue;
				pValue = new TCHAR[dwSize + 1];

				queryValue(pszValueName, nValueType, (LPBYTE)pValue, dwSize, moreData);
			}*/
			if(moreData)
			{
				strValue = _T("");
				return;
			}

			if (REG_EXPAND_SZ == nValueType && bAutoExpand)
			{
				TCHAR* pNewValue = new TCHAR[dwSize];
				DWORD dwNewSize = ::ExpandEnvironmentStrings(pValue, pNewValue, dwSize);

				if (dwNewSize > dwSize)
				{
					delete [] pNewValue;

					pNewValue = new TCHAR[dwNewSize];

					::ExpandEnvironmentStrings(pNewValue, pValue, dwNewSize);
				}

				//delete [] pValue;
				strValue = pNewValue;
				delete [] pNewValue;
				return;
			}

			strValue = pValue;
			//delete [] pValue;
		}

		void queryValue(PCTSTR pszValueName, DWORD& dwValue) throw(_win32_error)
		{
			DWORD dwSize = sizeof(DWORD);
			DWORD nValueType = 0;
			bool moreData;

			queryValue(pszValueName, nValueType, (LPBYTE)&dwValue, dwSize, moreData);

			if (REG_DWORD != nValueType || REG_DWORD_LITTLE_ENDIAN != nValueType)
				throw _win32_error(ERROR_INVALID_FUNCTION);
		}

		// Boolean
		void setValue(PCTSTR pszValueName) throw(_win32_error)
		{
			setValue(pszValueName, REG_NONE, NULL, 0);
		}

		// Binary
		void setValue(PCTSTR pszValueName, PVOID pData, DWORD dwDataLen) throw(_win32_error)
		{
			setValue(pszValueName, REG_BINARY, (LPBYTE)pData, dwDataLen);
		}

		// String
		void setValue(PCTSTR pszValueName, PCTSTR pszValue) throw(_win32_error)
		{
			// Lengths of string types must include the null terminator
			DWORD nDataSize = ((DWORD)_tcslen(pszValue) + 1) * sizeof(TCHAR);
			setValue(pszValueName, REG_SZ, (LPBYTE)pszValue, nDataSize);
		}

		// DWORD
		void setValue(PCTSTR pszValueName, DWORD dwValue) throw(_win32_error)
		{
			setValue(pszValueName, REG_DWORD, (LPBYTE)&dwValue, sizeof(DWORD));
		}

		struct KeyNameMap
		{
			HKEY	hkey;
			TCHAR*	pszName;
		};

		static const KeyNameMap* GetKeyNameMap()
		{
			static KeyNameMap s_nameMap[] = 
			{
				HKEY_CURRENT_USER,		_T("HKEY_CURRENT_USER"),
				HKEY_LOCAL_MACHINE,		_T("HKEY_LOCAL_MACHINE"),
				HKEY_CLASSES_ROOT,		_T("HKEY_CLASSES_ROOT"),
				HKEY_USERS,				_T("HKEY_USERS"),
				HKEY_PERFORMANCE_DATA,	_T("HKEY_PERFORMANCE_DATA"),
				HKEY_CURRENT_CONFIG,	_T("HKEY_CURRENT_CONFIG"),
				HKEY_DYN_DATA,			_T("HKEY_DYN_DATA"),
				HKEY_CURRENT_USER,		_T("HKCU"),
				HKEY_LOCAL_MACHINE,		_T("HKLM"),
				HKEY_CLASSES_ROOT,		_T("HKCR"),
				HKEY_USERS,				_T("HKU"),
				HKEY_PERFORMANCE_DATA,	_T("HKPD"),
				HKEY_CURRENT_CONFIG,	_T("HKCC"),
				HKEY_DYN_DATA,			_T("HKDD"),
				NULL,					0
			};
			
			return s_nameMap;
		}

		static const TCHAR* getRootKeyName(HKEY key)
		{
			const KeyNameMap* s_nameMap = GetKeyNameMap();

			for (int i = 0; s_nameMap[i].hkey != NULL; i++)
			{
				if (s_nameMap[i].hkey == key)
					return s_nameMap[i].pszName;
			}

			return NULL;
		}

		static HKEY getRootKeyFromName(const TCHAR* pszName) throw()
		{
			const KeyNameMap* s_nameMap = GetKeyNameMap();
			
			for (int i = 0; s_nameMap[i].hkey != NULL; i++)
			{
				if (_tcscmp(pszName, s_nameMap[i].pszName) == 0)
					return s_nameMap[i].hkey;
			}
			
			return NULL;
		}
   
		static tstring stripFirstKeyName(const TCHAR* p) throw()
		{
			const TCHAR *pp = _tcschr(p, '\\');
		
			if (pp == NULL)
				return tstring(p);
			else
				return tstring(pp - p + 1, 0);
		}

	protected:
		// Hide the dangerous functions
		void queryValue(
			PCTSTR pszValueName, 
			DWORD& dwValueType,
			LPBYTE lpData, 
			DWORD& dwByteSize,
			bool& moreData) throw(_win32_error)
		{
			DWORD error = ::RegQueryValueEx(
				m_hkey,
				pszValueName,
				NULL,
				&dwValueType,
				lpData,
				&dwByteSize);
				
			if (error != ERROR_SUCCESS && error != ERROR_MORE_DATA)
				throw _win32_error(error);
				
			moreData = (error == ERROR_MORE_DATA);
		}

		void setValue(
			PCTSTR pszValueName, 
			DWORD nValueType, 
			LPBYTE lpData, 
			DWORD nDataSize) throw(_win32_error)
		{
			THROW_BAD_WINERROR_RET(::RegSetValueEx(
				m_hkey, pszValueName, NULL, 
				nValueType, lpData, nDataSize));
		}

		// Data
		HKEY m_hkey;
	};
}

#pragma warning(pop)
